#ifndef __TENSOR_H
#define __TENSOR_H

#include <vector>
#include <string>
#include "opencl_wrapper.h"
#include "enviroment.h"
#include "utils.h"
// #include "operation.h"

namespace mslite{

enum MemType {
    Buffer = 1,
    Image2d = 2
};

enum DataType {
    clint8 = 1,
    clfp32 = 2
};

enum Format {
    NCHW = 1,   // for 4D feature
    NCHW4 = 2,
    NHWC = 3,
    NHWC4 = 4,  

    OHWI = 5,   // for 4D weight
    OIHW = 6,
    OHWII4O4 = 7,

    HW = 8,     // for 2D data
    WH = 9,
};

template <typename Dtype> class Operation;
class Enviroment;

template <typename Dtype>
class Tensor {
friend class Operation<Dtype>;
public:
    Tensor() = default;
    Tensor(std::string name, std::vector<int> shape);
    Tensor(std::string name, std::vector<int> shape, Format dformat, MemType memtype);
    Tensor(std::string name, std::vector<int> shape, Format dformat, MemType memtype, const std::vector<Dtype>& data);
    int eleNums(void);
    int eleNums(std::vector<int> shape);
    int setData(std::vector<Dtype> data, Format srcFmt, int packTile);
    std::vector<Dtype> getData(void);
    std::vector<int>& shape(void) {return shape_;}
    cl_mem& memObj(void) {return memObj_;}
private: 
    std::string name_;
    std::vector<int> shape_;
    std::vector<Dtype> data_;
    Format dformat_;
    MemType memtype_;
    cl_mem memObj_;
};

template <typename Dtype>
Tensor<Dtype>::Tensor(std::string name, std::vector<int> shape) {
    MemType memtype = Buffer;
    Format dformat;
    if (shape.size() == 4)
        dformat = NCHW;
    if (shape.size() == 2)
        dformat = HW;
    Tensor(name, shape, dformat, memtype);
}

template <typename Dtype>
Tensor<Dtype>::Tensor(std::string name, std::vector<int> shape, Format dformat, MemType memtype):
    Tensor(name, shape, dformat, memtype, std::vector<Dtype>(eleNums(shape), 0)) {
    
}

template <typename Dtype>
Tensor<Dtype>::Tensor(std::string name, std::vector<int> shape, Format dformat, MemType memtype, const std::vector<Dtype>& data):
    name_(name), shape_(shape), dformat_(dformat), memtype_(memtype), data_(data) {
    cl_int status;
    // std::cout << name << ": " << eleNums() <<std::endl;
    if (memtype_ == Buffer) {
        memObj_ = clCreateBuffer(Enviroment::context_, CL_MEM_READ_WRITE |  CL_MEM_COPY_HOST_PTR, eleNums()*sizeof(Dtype), data_.data(), &status);

    }
    if (memtype_ == Image2d) {
        cl_channel_type channel_dtype = sizeof(Dtype) == 1 ? CL_SIGNED_INT8 : CL_FLOAT;
        memObj_ = createImage2d(Enviroment::context_, channel_dtype, shape_[1]/4, shape_[0], data_.data(), &status);
    }
}

template <typename Dtype>
int Tensor<Dtype>::eleNums(void) {
    if (shape_.size() == 0)
        return 0;
    int totalSize = 1;
    for (auto dimLen: shape_) {
        totalSize *= dimLen;
    }
    return totalSize;
}

template <typename Dtype>
int Tensor<Dtype>::eleNums(std::vector<int> shape) {
    if (shape.size() == 0)
        return 0;
    int totalSize = 1;
    for (auto dimLen: shape) {
        totalSize *= dimLen;
    }
    return totalSize;
}

template <typename Dtype>
int Tensor<Dtype>::setData(std::vector<Dtype> data, Format srcFmt, int packTile) {
    // std::cout << this->name_ << " " << data.size() << " " << eleNums() << std::endl;
    assert(data.size() == eleNums());
    cl_int status = 0;
    data_ = std::move(data);
    if (srcFmt != dformat_) {
        if (srcFmt == HW && dformat_ == WH) {
            data_ = reorderHWToWH(data_, shape_);
        }
    }
    std::vector<Dtype> packed = packInRow(data_, shape_, packTile);
    if (memtype_ == Buffer) {
        status = clEnqueueWriteBuffer(Enviroment::queue_, memObj_, true, 0, eleNums()*sizeof(Dtype), packed.data(), 0, nullptr, nullptr);
    }
    if (memtype_ == Image2d) {
        size_t width = shape_[1]/4, height = shape_[0];
        size_t origin[3] = {0, 0, 0};  // 0, 0, 0
        size_t region[3] = {width, height, 1};  // width_, height_, 1
        status = clEnqueueWriteImage(Enviroment::queue_, memObj_, CL_TRUE, origin, region, 0, 0, packed.data(), 0, NULL, NULL);
        CHECK(status);
    }
    return status;
}

template <typename Dtype>
std::vector<Dtype> Tensor<Dtype>::getData(void) {
    cl_int status;
    if (memtype_ == Buffer) {
        status = clEnqueueReadBuffer(Enviroment::queue_, memObj_, CL_TRUE, 0, sizeof(Dtype)*eleNums(), data_.data(), 0, NULL, NULL);
    }
    if (memtype_ == Image2d) {

    }
    return data_;
}

}

#endif